"""
Base MCP server implementation for CodeMCP framework.

Provides the foundational MCP server class that all language-specific
and document analysis MCPs inherit from.
"""

import asyncio
from typing import Dict, Any, List, Optional, Callable, Union
from abc import ABC, abstractmethod
import json
from datetime import datetime
from pathlib import Path

from .config import Config
from .tool_registry import ToolRegistry, ToolType, tool, get_tool_registry
from .error_handler import CodeMCPError, AnalysisError, handle_error, log_info, log_debug
from .monitoring import HealthMonitor


class CallTreeNode:
    """Represents a node in the function call tree."""
    
    def __init__(self,
                 name: str,
                 file_path: str,
                 line_number: int,
                 language: str,
                 node_type: str = "function",
                 signature: str = None,
                 documentation: str = None):
        self.name = name
        self.file_path = file_path
        self.line_number = line_number
        self.language = language
        self.node_type = node_type  # function, method, class, section, etc.
        self.signature = signature
        self.documentation = documentation
        
        # Call relationships
        self.caller_nodes: List['CallTreeNode'] = []
        self.callee_nodes: List['CallTreeNode'] = []
        
        # Metadata
        self.call_frequency = 0
        self.call_context = ""
        self.short_description = ""
        self.llm_explanation = ""
        self.semantic_tags: List[str] = []
        self.complexity_score = 0.0
        
        # Timestamps
        self.created_at = datetime.utcnow()
        self.updated_at = datetime.utcnow()
    
    def add_callee(self, node: 'CallTreeNode'):
        """Add a callee node (function this node calls)."""
        if node not in self.callee_nodes:
            self.callee_nodes.append(node)
        if self not in node.caller_nodes:
            node.caller_nodes.append(self)
    
    def add_caller(self, node: 'CallTreeNode'):
        """Add a caller node (function that calls this node)."""
        if node not in self.caller_nodes:
            self.caller_nodes.append(node)
        if self not in node.callee_nodes:
            node.callee_nodes.append(self)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert node to dictionary representation."""
        return {
            "name": self.name,
            "file_path": self.file_path,
            "line_number": self.line_number,
            "language": self.language,
            "node_type": self.node_type,
            "signature": self.signature,
            "documentation": self.documentation,
            "call_frequency": self.call_frequency,
            "call_context": self.call_context,
            "short_description": self.short_description,
            "llm_explanation": self.llm_explanation,
            "semantic_tags": self.semantic_tags,
            "complexity_score": self.complexity_score,
            "caller_count": len(self.caller_nodes),
            "callee_count": len(self.callee_nodes),
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat()
        }
    
    def __str__(self):
        return f"{self.name}@{self.file_path}:{self.line_number}"
    
    def __repr__(self):
        return f"CallTreeNode(name='{self.name}', file='{self.file_path}', line={self.line_number})"


class CallTree:
    """Represents a complete function call tree."""
    
    def __init__(self, root_node: CallTreeNode = None):
        self.root_node = root_node
        self.nodes: Dict[str, CallTreeNode] = {}
        self.metadata = {
            "created_at": datetime.utcnow(),
            "analysis_type": "unknown",
            "language": "unknown",
            "total_nodes": 0,
            "max_depth": 0
        }
    
    def add_node(self, node: CallTreeNode) -> str:
        """Add a node to the tree and return its ID."""
        node_id = f"{node.file_path}:{node.line_number}:{node.name}"
        self.nodes[node_id] = node
        self.metadata["total_nodes"] = len(self.nodes)
        return node_id
    
    def get_node(self, node_id: str) -> Optional[CallTreeNode]:
        """Get a node by ID."""
        return self.nodes.get(node_id)
    
    def calculate_depth(self) -> int:
        """Calculate maximum depth of the tree."""
        if not self.root_node:
            return 0
        
        def _get_depth(node: CallTreeNode, visited: set) -> int:
            if id(node) in visited:
                return 0  # Avoid infinite recursion
            
            visited.add(id(node))
            if not node.callee_nodes:
                return 1
            
            max_child_depth = max(_get_depth(child, visited.copy()) 
                                for child in node.callee_nodes)
            return 1 + max_child_depth
        
        depth = _get_depth(self.root_node, set())
        self.metadata["max_depth"] = depth
        return depth
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert tree to dictionary representation."""
        return {
            "root_node": self.root_node.to_dict() if self.root_node else None,
            "nodes": {node_id: node.to_dict() for node_id, node in self.nodes.items()},
            "metadata": self.metadata
        }


class BaseMCP(ABC):
    """
    Base MCP server class for CodeMCP framework.
    
    All language-specific and document analysis MCPs inherit from this class.
    """
    
    def __init__(self, 
                 name: str,
                 language: str = None,
                 config: Config = None):
        self.name = name
        self.language = language
        self.config = config or Config()
        
        # Initialize components
        self.tool_registry = get_tool_registry()
        self.health_monitor = HealthMonitor(self.name)
        
        # State management
        self.is_initialized = False
        self.is_running = False
        self.session_data: Dict[str, Any] = {}
        
        # Statistics
        self.stats = {
            "requests_processed": 0,
            "errors_encountered": 0,
            "avg_processing_time": 0.0,
            "last_activity": None
        }
        
        log_info(f"Initialized {self.name} MCP server")
    
    async def initialize(self):
        """Initialize the MCP server."""
        try:
            # Register default tools
            await self._register_default_tools()
            
            # Initialize language-specific components
            await self._initialize_language_components()
            
            # Start health monitoring
            await self.health_monitor.start()
            
            self.is_initialized = True
            log_info(f"{self.name} MCP server initialized successfully")
            
        except Exception as e:
            error_response = handle_error(e, {"mcp_server": self.name})
            raise CodeMCPError(f"Failed to initialize {self.name}: {error_response['message']}")
    
    async def start(self):
        """Start the MCP server."""
        if not self.is_initialized:
            await self.initialize()
        
        self.is_running = True
        log_info(f"{self.name} MCP server started")
    
    async def stop(self):
        """Stop the MCP server."""
        self.is_running = False
        await self.health_monitor.stop()
        log_info(f"{self.name} MCP server stopped")
    
    @abstractmethod
    async def _register_default_tools(self):
        """Register default tools for this MCP server."""
        pass
    
    @abstractmethod
    async def _initialize_language_components(self):
        """Initialize language-specific components."""
        pass
    
    @abstractmethod
    async def parse_file(self, file_path: str) -> Any:
        """Parse a file and return AST or equivalent structure."""
        pass
    
    @abstractmethod
    async def build_call_tree(self, 
                             entry_point: str, 
                             max_depth: int = None) -> CallTree:
        """Build a call tree starting from entry point."""
        pass
    
    # Common MCP tools that all servers should have
    
    async def get_server_info(self) -> Dict[str, Any]:
        """Get information about this MCP server."""
        return {
            "name": self.name,
            "language": self.language,
            "version": "1.0.0",
            "status": "running" if self.is_running else "stopped",
            "initialized": self.is_initialized,
            "registered_tools": len(self.tool_registry.list_tools()),
            "statistics": self.stats,
            "health": await self.health_monitor.get_health_status()
        }
    
    async def list_available_tools(self) -> List[Dict[str, Any]]:
        """List all available tools for this MCP server."""
        # Get tools for this language
        language_tools = self.tool_registry.list_tools(language=self.language)
        
        # Get general tools
        general_tools = self.tool_registry.list_tools(language=None)
        
        all_tools = language_tools + general_tools
        return [tool.to_dict() for tool in all_tools]
    
    async def validate_file(self, file_path: str) -> Dict[str, Any]:
        """Validate if a file can be processed by this MCP server."""
        file_path = Path(file_path)
        
        # Check if file exists
        if not file_path.exists():
            return {
                "valid": False,
                "reason": "File does not exist",
                "file_path": str(file_path)
            }
        
        # Check file extension
        if not self._is_supported_file(file_path):
            return {
                "valid": False,
                "reason": f"File extension not supported by {self.name}",
                "file_path": str(file_path),
                "supported_extensions": self._get_supported_extensions()
            }
        
        # Try basic parsing
        try:
            await self.parse_file(str(file_path))
            return {
                "valid": True,
                "file_path": str(file_path),
                "language": self.language,
                "file_size": file_path.stat().st_size
            }
        except Exception as e:
            return {
                "valid": False,
                "reason": f"Parse error: {str(e)}",
                "file_path": str(file_path)
            }
    
    @abstractmethod
    def _is_supported_file(self, file_path: Path) -> bool:
        """Check if file is supported by this MCP server."""
        pass
    
    @abstractmethod
    def _get_supported_extensions(self) -> List[str]:
        """Get list of supported file extensions."""
        pass
    
    async def get_function_info(self, 
                               file_path: str, 
                               function_name: str) -> Dict[str, Any]:
        """Get detailed information about a specific function."""
        try:
            # Parse the file
            ast_data = await self.parse_file(file_path)
            
            # Extract function information
            function_info = await self._extract_function_info(ast_data, function_name)
            
            if not function_info:
                raise AnalysisError(f"Function '{function_name}' not found in {file_path}")
            
            return function_info
            
        except Exception as e:
            error_response = handle_error(e, {
                "file_path": file_path,
                "function_name": function_name,
                "mcp_server": self.name
            })
            raise AnalysisError(f"Failed to get function info: {error_response['message']}")
    
    @abstractmethod
    async def _extract_function_info(self, 
                                   ast_data: Any, 
                                   function_name: str) -> Optional[Dict[str, Any]]:
        """Extract function information from AST data."""
        pass
    
    async def analyze_dependencies(self, file_path: str) -> Dict[str, Any]:
        """Analyze file dependencies (imports, includes, etc.)."""
        try:
            # Parse the file
            ast_data = await self.parse_file(file_path)
            
            # Extract dependencies
            dependencies = await self._extract_dependencies(ast_data)
            
            return {
                "file_path": file_path,
                "dependencies": dependencies,
                "dependency_count": len(dependencies),
                "analysis_timestamp": datetime.utcnow().isoformat()
            }
            
        except Exception as e:
            error_response = handle_error(e, {
                "file_path": file_path,
                "mcp_server": self.name
            })
            raise AnalysisError(f"Failed to analyze dependencies: {error_response['message']}")
    
    @abstractmethod
    async def _extract_dependencies(self, ast_data: Any) -> List[Dict[str, Any]]:
        """Extract dependencies from AST data."""
        pass
    
    async def get_call_chain(self, 
                           start_function: str,
                           end_function: str,
                           file_path: str = None,
                           max_depth: int = None) -> Dict[str, Any]:
        """Find call chain between two functions."""
        try:
            # Build call tree
            if file_path:
                call_tree = await self.build_call_tree(f"{file_path}::{start_function}", max_depth)
            else:
                call_tree = await self.build_call_tree(start_function, max_depth)
            
            # Find path from start to end function
            path = self._find_call_path(call_tree, start_function, end_function)
            
            return {
                "start_function": start_function,
                "end_function": end_function,
                "path_found": path is not None,
                "call_path": path,
                "path_length": len(path) if path else 0,
                "analysis_timestamp": datetime.utcnow().isoformat()
            }
            
        except Exception as e:
            error_response = handle_error(e, {
                "start_function": start_function,
                "end_function": end_function,
                "file_path": file_path,
                "mcp_server": self.name
            })
            raise AnalysisError(f"Failed to find call chain: {error_response['message']}")
    
    def _find_call_path(self, 
                       call_tree: CallTree, 
                       start_function: str, 
                       end_function: str) -> Optional[List[str]]:
        """Find path between two functions in call tree."""
        # Simple BFS implementation
        if not call_tree.root_node:
            return None
        
        # Find start node
        start_node = None
        for node in call_tree.nodes.values():
            if node.name == start_function:
                start_node = node
                break
        
        if not start_node:
            return None
        
        # BFS to find end function
        queue = [(start_node, [start_function])]
        visited = set()
        
        while queue:
            current_node, path = queue.pop(0)
            
            if id(current_node) in visited:
                continue
            visited.add(id(current_node))
            
            if current_node.name == end_function:
                return path
            
            # Add callee nodes to queue
            for callee in current_node.callee_nodes:
                if id(callee) not in visited:
                    queue.append((callee, path + [callee.name]))
        
        return None
    
    async def process_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """Process an incoming MCP request."""
        start_time = datetime.utcnow()
        
        try:
            # Update statistics
            self.stats["requests_processed"] += 1
            self.stats["last_activity"] = start_time.isoformat()
            
            # Validate request
            if "tool" not in request:
                raise CodeMCPError("Missing 'tool' field in request")
            
            tool_name = request["tool"]
            parameters = request.get("parameters", {})
            
            # Call the tool
            result = await self.tool_registry.call_tool(tool_name, **parameters)
            
            # Calculate processing time
            end_time = datetime.utcnow()
            processing_time = (end_time - start_time).total_seconds()
            
            # Update average processing time
            old_avg = self.stats["avg_processing_time"]
            count = self.stats["requests_processed"]
            self.stats["avg_processing_time"] = (old_avg * (count - 1) + processing_time) / count
            
            return {
                "success": True,
                "result": result,
                "processing_time": processing_time,
                "timestamp": end_time.isoformat()
            }
            
        except Exception as e:
            # Update error statistics
            self.stats["errors_encountered"] += 1
            
            error_response = handle_error(e, {
                "request": request,
                "mcp_server": self.name
            })
            
            return {
                "success": False,
                "error": error_response,
                "timestamp": datetime.utcnow().isoformat()
            }
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get server statistics."""
        return {
            "server_name": self.name,
            "language": self.language,
            "statistics": self.stats,
            "health": self.health_monitor.get_current_health(),
            "tool_registry": self.tool_registry.get_registry_info()
        }